use crate::transformer::{Transformer, TransformerConfig};
use crate::nlp::Tokenizer;
use rand::Rng;
use serde::{Deserialize, Serialize};

/// 生成式语言模型
#[derive(Clone, Serialize, Deserialize)]
pub struct LanguageModel {
    transformer: Transformer,
    tokenizer: Tokenizer,
    vocab_size: usize,
    embedding_dim: usize,
}

impl LanguageModel {
    /// 创建新的语言模型
    pub fn new(config: TransformerConfig, tokenizer: Tokenizer) -> Self {
        let vocab_size = tokenizer.vocab_size();
        let embedding_dim = config.d_model;
        
        LanguageModel {
            transformer: Transformer::new(config),
            tokenizer,
            vocab_size,
            embedding_dim,
        }
    }

    /// 将token IDs转换为嵌入
    fn ids_to_embeddings(&self, ids: &[usize]) -> Vec<Vec<f64>> {
        ids.iter()
            .map(|&id| {
                let mut embedding = vec![0.0; self.embedding_dim];
                // 简单的one-hot编码 + 缩放
                let idx = id % self.embedding_dim;
                embedding[idx] = 1.0;
                
                // 添加一些随机性使嵌入更丰富
                let scale = 0.1;
                let mut rng = rand::thread_rng();
                for i in 0..self.embedding_dim {
                    if i != idx {
                        embedding[i] = rng.gen_range(-scale..scale);
                    }
                }
                embedding
            })
            .collect()
    }

    /// 前向传播：输入文本 -> 下一个token的概率分布
    pub fn forward(&self, text: &str) -> Vec<f64> {
        // 分词
        let tokens = self.tokenizer.tokenize(text);
        let ids = self.tokenizer.encode(&tokens);
        
        // 转换为嵌入
        let embeddings = self.ids_to_embeddings(&ids);
        
        // Transformer处理
        let output = self.transformer.forward(&embeddings);
        
        // 使用最后一个token的输出
        if let Some(last_output) = output.last() {
            // 投影到词表大小
            self.project_to_vocab(last_output)
        } else {
            vec![0.0; self.vocab_size]
        }
    }

    /// 将Transformer输出投影到词表大小
    fn project_to_vocab(&self, hidden: &[f64]) -> Vec<f64> {
        // 简化：使用hidden state的不同维度组合来预测每个token
        let mut logits = vec![0.0; self.vocab_size];
        
        for i in 0..self.vocab_size.min(hidden.len()) {
            logits[i] = hidden[i % hidden.len()];
        }
        
        // Softmax归一化
        softmax(&logits)
    }

    /// 训练一步
    pub fn train_step(&mut self, input_text: &str, target_text: &str, learning_rate: f64) -> f64 {
        // 编码输入和目标
        let input_tokens = self.tokenizer.tokenize(input_text);
        let target_tokens = self.tokenizer.tokenize(target_text);
        
        let input_ids = self.tokenizer.encode(&input_tokens);
        let target_ids = self.tokenizer.encode(&target_tokens);
        
        // 转换为嵌入
        let input_embeddings = self.ids_to_embeddings(&input_ids);
        let target_embeddings = self.ids_to_embeddings(&target_ids);
        
        // 调用Transformer训练
        self.transformer.train_step(&input_embeddings, &target_embeddings, learning_rate)
    }

    /// 生成文本
    pub fn generate(
        &self,
        prompt: &str,
        max_length: usize,
        temperature: f64,
        top_k: Option<usize>,
        top_p: Option<f64>,
    ) -> String {
        let mut current_text = prompt.to_string();
        let mut generated_tokens = Vec::new();

        for _ in 0..max_length {
            // 获取下一个token的概率分布
            let probs = self.forward(&current_text);
            
            // 采样下一个token
            let next_id = self.sample_token(&probs, temperature, top_k, top_p);
            
            // 解码token
            let next_tokens = self.tokenizer.decode(&[next_id]);
            if next_tokens.is_empty() {
                break;
            }
            
            let next_token = &next_tokens[0];
            
            // 停止条件
            if next_token == "[SEP]" || next_token == "。" || next_token == "！" {
                generated_tokens.push(next_token.clone());
                break;
            }
            
            generated_tokens.push(next_token.clone());
            current_text.push_str(next_token);
        }

        generated_tokens.join("")
    }

    /// 采样下一个token
    fn sample_token(
        &self,
        probs: &[f64],
        temperature: f64,
        top_k: Option<usize>,
        top_p: Option<f64>,
    ) -> usize {
        // 应用温度
        let mut adjusted_probs: Vec<f64> = probs
            .iter()
            .map(|&p| (p.ln() / temperature).exp())
            .collect();
        
        // 归一化
        let sum: f64 = adjusted_probs.iter().sum();
        if sum > 0.0 {
            for p in &mut adjusted_probs {
                *p /= sum;
            }
        }

        // Top-k采样
        if let Some(k) = top_k {
            let mut indexed_probs: Vec<(usize, f64)> = adjusted_probs
                .iter()
                .enumerate()
                .map(|(i, &p)| (i, p))
                .collect();
            
            indexed_probs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
            
            let top_k_probs: Vec<(usize, f64)> = indexed_probs.into_iter().take(k).collect();
            
            // 重新归一化
            let sum: f64 = top_k_probs.iter().map(|(_, p)| p).sum();
            if sum > 0.0 {
                adjusted_probs = vec![0.0; probs.len()];
                for (idx, prob) in top_k_probs {
                    adjusted_probs[idx] = prob / sum;
                }
            }
        }

        // Top-p (nucleus) 采样
        if let Some(p) = top_p {
            let mut indexed_probs: Vec<(usize, f64)> = adjusted_probs
                .iter()
                .enumerate()
                .map(|(i, &prob)| (i, prob))
                .collect();
            
            indexed_probs.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
            
            let mut cumsum = 0.0;
            let mut nucleus = Vec::new();
            
            for (idx, prob) in indexed_probs {
                cumsum += prob;
                nucleus.push((idx, prob));
                if cumsum >= p {
                    break;
                }
            }
            
            // 重新归一化
            let sum: f64 = nucleus.iter().map(|(_, p)| p).sum();
            if sum > 0.0 {
                adjusted_probs = vec![0.0; probs.len()];
                for (idx, prob) in nucleus {
                    adjusted_probs[idx] = prob / sum;
                }
            }
        }

        // 从分布中采样
        let mut rng = rand::thread_rng();
        let rand_val: f64 = rng.gen();
        let mut cumsum = 0.0;
        
        for (i, &prob) in adjusted_probs.iter().enumerate() {
            cumsum += prob;
            if rand_val < cumsum {
                return i;
            }
        }
        
        // 默认返回最后一个
        adjusted_probs.len() - 1
    }

    /// 获取tokenizer引用
    pub fn tokenizer(&self) -> &Tokenizer {
        &self.tokenizer
    }

    /// 获取transformer引用
    pub fn transformer(&self) -> &Transformer {
        &self.transformer
    }
}

/// Softmax函数
fn softmax(logits: &[f64]) -> Vec<f64> {
    let max_logit = logits.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
    
    let exp_logits: Vec<f64> = logits.iter().map(|&x| (x - max_logit).exp()).collect();
    let sum: f64 = exp_logits.iter().sum();
    
    if sum > 0.0 {
        exp_logits.iter().map(|&x| x / sum).collect()
    } else {
        vec![1.0 / logits.len() as f64; logits.len()]
    }
}

/// 对话数据集
#[derive(Debug, Clone)]
pub struct DialogueDataset {
    pub pairs: Vec<(String, String)>, // (输入, 回复)
}

impl DialogueDataset {
    /// 创建新的数据集
    pub fn new() -> Self {
        DialogueDataset { pairs: Vec::new() }
    }

    /// 添加对话对
    pub fn add_pair(&mut self, input: String, response: String) {
        self.pairs.push((input, response));
    }

    /// 从示例创建数据集
    pub fn from_examples() -> Self {
        let mut dataset = Self::new();
        
        // 添加基础对话示例
        dataset.add_pair("你好".to_string(), "你好！很高兴见到你。".to_string());
        dataset.add_pair("你好吗".to_string(), "我很好，谢谢你的关心！你呢？".to_string());
        dataset.add_pair("你是谁".to_string(), "我是一个神经符号AI系统，可以帮助你解决问题。".to_string());
        dataset.add_pair("你能做什么".to_string(), "我可以进行对话、回答问题、帮你理解各种概念。".to_string());
        
        // 技术相关
        dataset.add_pair("什么是AI".to_string(), "AI是人工智能，让机器能够模拟人类智能的技术。".to_string());
        dataset.add_pair("什么是机器学习".to_string(), "机器学习是AI的一个分支，通过数据让机器自动学习规律。".to_string());
        dataset.add_pair("什么是深度学习".to_string(), "深度学习使用多层神经网络来学习复杂的数据表示。".to_string());
        
        // 日常对话
        dataset.add_pair("今天天气真好".to_string(), "是啊，好天气总让人心情愉悦！".to_string());
        dataset.add_pair("谢谢你".to_string(), "不客气！很高兴能帮到你。".to_string());
        dataset.add_pair("再见".to_string(), "再见！期待下次和你聊天。".to_string());
        
        // 更多技术对话
        dataset.add_pair("Transformer是什么".to_string(), "Transformer是一种强大的神经网络架构，使用注意力机制处理序列。".to_string());
        dataset.add_pair("如何训练模型".to_string(), "需要准备数据、设计网络、选择优化器，然后迭代训练。".to_string());
        
        dataset
    }

    /// 获取数据集大小
    pub fn len(&self) -> usize {
        self.pairs.len()
    }

    /// 检查是否为空
    pub fn is_empty(&self) -> bool {
        self.pairs.is_empty()
    }
}

impl Default for DialogueDataset {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_softmax() {
        let logits = vec![1.0, 2.0, 3.0];
        let probs = softmax(&logits);
        
        let sum: f64 = probs.iter().sum();
        assert!((sum - 1.0).abs() < 1e-6);
        assert!(probs[2] > probs[1]);
        assert!(probs[1] > probs[0]);
    }

    #[test]
    fn test_dialogue_dataset() {
        let dataset = DialogueDataset::from_examples();
        assert!(!dataset.is_empty());
        assert!(dataset.len() > 5);
    }
}

